#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
插件系统基类模块
定义插件开发的标准接口和规范

插件开发规范:
1. 所有插件必须继承BasePlugin基类
2. 实现必需的接口方法: init(), run(), clean()
3. 提供插件元数据信息
4. 遵循安全和清理原则
"""

import os
import sys
import abc
import time
import json
import logging
import threading
from typing import Dict, List, Optional, Any, Union, Callable
from dataclasses import dataclass, field
from enum import Enum
import platform
import traceback

class PluginType(Enum):
    """插件类型枚举"""
    PERSISTENCE = "persistence"
    EVASION = "evasion"
    PRIVILEGE_ESCALATION = "privilege_escalation"
    DATA_EXFILTRATION = "data_exfiltration"
    LATERAL_MOVEMENT = "lateral_movement"
    RECONNAISSANCE = "reconnaissance"
    BACKDOOR = "backdoor"
    UTILITY = "utility"

class PluginStatus(Enum):
    """插件状态枚举"""
    UNINITIALIZED = "uninitialized"
    INITIALIZED = "initialized"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"
    CLEANED = "cleaned"

class PluginPriority(Enum):
    """插件优先级枚举"""
    LOW = 1
    MEDIUM = 2
    HIGH = 3
    CRITICAL = 4

@dataclass
class PluginMetadata:
    """插件元数据"""
    name: str
    version: str
    description: str
    author: str
    plugin_type: PluginType
    priority: PluginPriority = PluginPriority.MEDIUM
    dependencies: List[str] = field(default_factory=list)
    supported_platforms: List[str] = field(default_factory=lambda: ["linux"])
    requires_root: bool = False
    stealth_level: int = 3  # 1-5, 5为最隐蔽
    persistence_level: int = 3  # 1-5, 5为最持久
    cleanup_required: bool = True
    tags: List[str] = field(default_factory=list)

@dataclass
class PluginConfig:
    """插件配置"""
    enabled: bool = True
    auto_start: bool = False
    max_execution_time: int = 300  # 秒
    retry_count: int = 3
    cleanup_on_exit: bool = True
    log_level: str = "INFO"
    custom_params: Dict[str, Any] = field(default_factory=dict)

@dataclass
class PluginResult:
    """插件执行结果"""
    success: bool
    message: str
    data: Dict[str, Any] = field(default_factory=dict)
    execution_time: float = 0.0
    error: Optional[str] = None
    artifacts: List[str] = field(default_factory=list)  # 创建的文件/进程等

class BasePlugin(abc.ABC):
    """插件基类 - 所有插件必须继承此类"""
    
    def __init__(self, config: PluginConfig = None):
        """
        初始化插件
        
        Args:
            config: 插件配置
        """
        self.config = config or PluginConfig()
        self.status = PluginStatus.UNINITIALIZED
        self.metadata = self.get_metadata()
        self.logger = self._setup_logger()
        self.start_time = 0.0
        self.artifacts = []  # 跟踪创建的文件、进程等
        self.cleanup_callbacks = []  # 清理回调函数
        
        # 验证平台兼容性
        if not self._check_platform_compatibility():
            raise RuntimeError(f"Plugin {self.metadata.name} is not compatible with {platform.system()}")
        
        # 验证权限要求
        if self.metadata.requires_root and os.geteuid() != 0:
            raise RuntimeError(f"Plugin {self.metadata.name} requires root privileges")
    
    def _setup_logger(self) -> logging.Logger:
        """设置日志记录器"""
        logger = logging.getLogger(f"plugin.{self.metadata.name}")
        logger.setLevel(getattr(logging, self.config.log_level))
        
        if not logger.handlers:
            handler = logging.StreamHandler()
            formatter = logging.Formatter(
                f'[%(asctime)s] [PLUGIN:{self.metadata.name}] [%(levelname)s] %(message)s'
            )
            handler.setFormatter(formatter)
            logger.addHandler(handler)
        
        return logger
    
    def _check_platform_compatibility(self) -> bool:
        """检查平台兼容性"""
        current_platform = platform.system().lower()
        return current_platform in [p.lower() for p in self.metadata.supported_platforms]
    
    @abc.abstractmethod
    def get_metadata(self) -> PluginMetadata:
        """
        获取插件元数据 - 必须实现
        
        Returns:
            PluginMetadata: 插件元数据
        """
        pass
    
    @abc.abstractmethod
    def init(self) -> bool:
        """
        初始化插件 - 必须实现
        
        Returns:
            bool: 初始化是否成功
        """
        pass
    
    @abc.abstractmethod
    def run(self) -> PluginResult:
        """
        运行插件主逻辑 - 必须实现
        
        Returns:
            PluginResult: 执行结果
        """
        pass
    
    @abc.abstractmethod
    def clean(self) -> bool:
        """
        清理插件资源 - 必须实现
        
        Returns:
            bool: 清理是否成功
        """
        pass
    
    def validate_dependencies(self) -> bool:
        """验证插件依赖"""
        for dep in self.metadata.dependencies:
            try:
                __import__(dep)
            except ImportError:
                self.logger.error(f"Missing dependency: {dep}")
                return False
        return True
    
    def register_artifact(self, artifact_path: str):
        """注册需要清理的文件或资源"""
        self.artifacts.append(artifact_path)
        self.logger.debug(f"Registered artifact: {artifact_path}")
    
    def register_cleanup_callback(self, callback: Callable):
        """注册清理回调函数"""
        self.cleanup_callbacks.append(callback)
        self.logger.debug(f"Registered cleanup callback: {callback.__name__}")
    
    def execute(self) -> PluginResult:
        """
        执行插件完整流程
        
        Returns:
            PluginResult: 执行结果
        """
        self.start_time = time.time()
        
        try:
            # 1. 验证依赖
            if not self.validate_dependencies():
                return PluginResult(
                    success=False,
                    message="Dependency validation failed",
                    error="Missing required dependencies"
                )
            
            # 2. 初始化
            self.logger.info(f"Initializing plugin {self.metadata.name}")
            if not self.init():
                self.status = PluginStatus.FAILED
                return PluginResult(
                    success=False,
                    message="Plugin initialization failed",
                    error="init() method returned False"
                )
            
            self.status = PluginStatus.INITIALIZED
            self.logger.info("Plugin initialized successfully")
            
            # 3. 运行主逻辑
            self.logger.info("Starting plugin execution")
            self.status = PluginStatus.RUNNING
            
            result = self.run()
            result.execution_time = time.time() - self.start_time
            result.artifacts = self.artifacts.copy()
            
            if result.success:
                self.status = PluginStatus.COMPLETED
                self.logger.info(f"Plugin completed successfully in {result.execution_time:.2f}s")
            else:
                self.status = PluginStatus.FAILED
                self.logger.error(f"Plugin execution failed: {result.error}")
            
            return result
            
        except Exception as e:
            self.status = PluginStatus.FAILED
            error_msg = f"Plugin execution error: {str(e)}"
            self.logger.error(error_msg)
            self.logger.debug(traceback.format_exc())
            
            return PluginResult(
                success=False,
                message="Plugin execution failed with exception",
                error=error_msg,
                execution_time=time.time() - self.start_time,
                artifacts=self.artifacts.copy()
            )
        
        finally:
            # 4. 自动清理（如果配置启用）
            if self.config.cleanup_on_exit:
                self.cleanup()
    
    def cleanup(self) -> bool:
        """执行完整清理流程"""
        try:
            self.logger.info("Starting plugin cleanup")
            
            # 1. 执行插件自定义清理
            plugin_cleanup_success = self.clean()
            
            # 2. 执行注册的清理回调
            for callback in self.cleanup_callbacks:
                try:
                    callback()
                    self.logger.debug(f"Executed cleanup callback: {callback.__name__}")
                except Exception as e:
                    self.logger.error(f"Cleanup callback failed: {e}")
            
            # 3. 清理注册的文件和资源
            for artifact in self.artifacts:
                try:
                    if os.path.exists(artifact):
                        if os.path.isfile(artifact):
                            os.remove(artifact)
                        elif os.path.isdir(artifact):
                            import shutil
                            shutil.rmtree(artifact)
                        self.logger.debug(f"Cleaned artifact: {artifact}")
                except Exception as e:
                    self.logger.error(f"Failed to clean artifact {artifact}: {e}")
            
            self.status = PluginStatus.CLEANED
            self.logger.info("Plugin cleanup completed")
            return plugin_cleanup_success
            
        except Exception as e:
            self.logger.error(f"Cleanup failed: {e}")
            return False
    
    def get_status(self) -> Dict[str, Any]:
        """获取插件状态信息"""
        return {
            "name": self.metadata.name,
            "version": self.metadata.version,
            "status": self.status.value,
            "type": self.metadata.plugin_type.value,
            "priority": self.metadata.priority.value,
            "execution_time": time.time() - self.start_time if self.start_time > 0 else 0,
            "artifacts_count": len(self.artifacts),
            "platform": platform.system(),
            "requires_root": self.metadata.requires_root,
            "stealth_level": self.metadata.stealth_level
        }
    
    def is_compatible(self) -> bool:
        """检查插件是否与当前环境兼容"""
        return (
            self._check_platform_compatibility() and
            (not self.metadata.requires_root or os.geteuid() == 0) and
            self.validate_dependencies()
        )

class PluginManager:
    """插件管理器"""
    
    def __init__(self, plugins_dir: str = None):
        """
        初始化插件管理器
        
        Args:
            plugins_dir: 插件目录路径
        """
        if plugins_dir is None:
            plugins_dir = os.path.join(os.path.dirname(__file__), "plugins")
        
        self.plugins_dir = plugins_dir
        self.loaded_plugins = {}
        self.plugin_configs = {}
        self.execution_history = []
        self.logger = logging.getLogger("plugin_manager")
    
    def discover_plugins(self) -> List[str]:
        """发现可用的插件"""
        plugins = []
        
        if not os.path.exists(self.plugins_dir):
            return plugins
        
        for item in os.listdir(self.plugins_dir):
            plugin_path = os.path.join(self.plugins_dir, item)
            
            # 检查Python文件
            if item.endswith('.py') and not item.startswith('__'):
                plugins.append(item[:-3])  # 移除.py扩展名
            
            # 检查插件目录
            elif os.path.isdir(plugin_path):
                init_file = os.path.join(plugin_path, '__init__.py')
                if os.path.exists(init_file):
                    plugins.append(item)
        
        return plugins
    
    def load_plugin(self, plugin_name: str, config: PluginConfig = None) -> bool:
        """
        加载插件
        
        Args:
            plugin_name: 插件名称
            config: 插件配置
            
        Returns:
            bool: 加载是否成功
        """
        try:
            # 动态导入插件模块
            if self.plugins_dir not in sys.path:
                sys.path.insert(0, self.plugins_dir)
            
            module = __import__(plugin_name)
            
            # 查找插件类（继承自BasePlugin的类）
            plugin_class = None
            for attr_name in dir(module):
                attr = getattr(module, attr_name)
                if (isinstance(attr, type) and 
                    issubclass(attr, BasePlugin) and 
                    attr != BasePlugin):
                    plugin_class = attr
                    break
            
            if not plugin_class:
                self.logger.error(f"No plugin class found in {plugin_name}")
                return False
            
            # 创建插件实例
            plugin_instance = plugin_class(config)
            
            # 检查兼容性
            if not plugin_instance.is_compatible():
                self.logger.error(f"Plugin {plugin_name} is not compatible with current environment")
                return False
            
            self.loaded_plugins[plugin_name] = plugin_instance
            self.plugin_configs[plugin_name] = config or PluginConfig()
            
            self.logger.info(f"Plugin {plugin_name} loaded successfully")
            return True
            
        except Exception as e:
            self.logger.error(f"Failed to load plugin {plugin_name}: {e}")
            return False
    
    def execute_plugin(self, plugin_name: str) -> Optional[PluginResult]:
        """
        执行插件
        
        Args:
            plugin_name: 插件名称
            
        Returns:
            PluginResult: 执行结果
        """
        if plugin_name not in self.loaded_plugins:
            self.logger.error(f"Plugin {plugin_name} not loaded")
            return None
        
        plugin = self.loaded_plugins[plugin_name]
        
        try:
            result = plugin.execute()
            
            # 记录执行历史
            self.execution_history.append({
                "plugin_name": plugin_name,
                "timestamp": time.time(),
                "success": result.success,
                "execution_time": result.execution_time,
                "message": result.message
            })
            
            return result
            
        except Exception as e:
            self.logger.error(f"Plugin execution failed: {e}")
            return PluginResult(
                success=False,
                message="Plugin execution failed with exception",
                error=str(e)
            )
    
    def get_plugin_info(self, plugin_name: str) -> Optional[Dict[str, Any]]:
        """获取插件信息"""
        if plugin_name not in self.loaded_plugins:
            return None
        
        plugin = self.loaded_plugins[plugin_name]
        return {
            "metadata": {
                "name": plugin.metadata.name,
                "version": plugin.metadata.version,
                "description": plugin.metadata.description,
                "author": plugin.metadata.author,
                "type": plugin.metadata.plugin_type.value,
                "priority": plugin.metadata.priority.value,
                "dependencies": plugin.metadata.dependencies,
                "supported_platforms": plugin.metadata.supported_platforms,
                "requires_root": plugin.metadata.requires_root,
                "stealth_level": plugin.metadata.stealth_level,
                "persistence_level": plugin.metadata.persistence_level,
                "tags": plugin.metadata.tags
            },
            "status": plugin.get_status(),
            "config": {
                "enabled": self.plugin_configs[plugin_name].enabled,
                "auto_start": self.plugin_configs[plugin_name].auto_start,
                "max_execution_time": self.plugin_configs[plugin_name].max_execution_time,
                "cleanup_on_exit": self.plugin_configs[plugin_name].cleanup_on_exit
            }
        }
    
    def list_plugins(self) -> Dict[str, Dict[str, Any]]:
        """列出所有已加载的插件"""
        return {name: self.get_plugin_info(name) for name in self.loaded_plugins.keys()}
    
    def cleanup_all_plugins(self):
        """清理所有插件"""
        for plugin_name, plugin in self.loaded_plugins.items():
            try:
                plugin.cleanup()
                self.logger.info(f"Cleaned up plugin: {plugin_name}")
            except Exception as e:
                self.logger.error(f"Failed to cleanup plugin {plugin_name}: {e}")

# 使用示例
if __name__ == "__main__":
    print("=== 插件系统基类测试 ===")
    
    # 这里只是展示基类的使用方式
    # 实际插件需要继承BasePlugin并实现抽象方法
    
    class ExamplePlugin(BasePlugin):
        def get_metadata(self) -> PluginMetadata:
            return PluginMetadata(
                name="example_plugin",
                version="1.0.0",
                description="示例插件",
                author="RedTeam",
                plugin_type=PluginType.UTILITY,
                supported_platforms=["linux", "windows"],
                tags=["example", "test"]
            )
        
        def init(self) -> bool:
            self.logger.info("Example plugin initialized")
            return True
        
        def run(self) -> PluginResult:
            self.logger.info("Example plugin running")
            return PluginResult(
                success=True,
                message="Example plugin executed successfully",
                data={"result": "Hello from example plugin!"}
            )
        
        def clean(self) -> bool:
            self.logger.info("Example plugin cleaned")
            return True
    
    # 测试插件
    config = PluginConfig(log_level="DEBUG")
    plugin = ExamplePlugin(config)
    result = plugin.execute()
    
    print(f"Plugin execution result: {result.success}")
    print(f"Message: {result.message}")
    print(f"Data: {result.data}")
    print(f"Execution time: {result.execution_time:.2f}s")